package org.example;

import java.util.*;
import static java.lang.Math.pow;
import static java.lang.Math.sqrt;

public class AStar {
    private final QuadTree root;

    AStar(QuadTree root) {
        this.root = root;
    }

    private static class PathNode implements Comparable<PathNode> {
        public Leaf treeNode;
        public PathNode parent;
        public double g;
        public double h;

        PathNode(Leaf treeNode, PathNode parent, double g, double h) {
            this.treeNode = treeNode;
            this.parent = parent;
            this.g = g;
            this.h = h;
        }

        public double f() {
            return g + h;
        }

        @Override
        public int compareTo(PathNode o) {
            return Double.compare(this.f(), o.f());
        }
    }

    private double distance(Leaf a, Leaf b) {
        var ac = a.center();
        var bc = b.center();
        return sqrt(pow(ac.x - bc.x, 2) + pow(ac.y - bc.y, 2));
    }

    // 在node的范围内查找坐标(x,y)所在的Leaf
    private Leaf searchRecursively(QuadTree node, int x, int y) {
        if (node.getClass() == Leaf.class) {
            Bounds bound = ((Leaf) node).bound;
            if (bound.contains(new Point(x, y)))
                return (Leaf)node;
            else
                return null;
        } else if(node.getClass() == TreeNode.class) {
            var leafOfNW = searchRecursively(((TreeNode) node).nw, x, y);
            if (leafOfNW != null) {
                return leafOfNW;
            }
            var leafOfNE = searchRecursively(((TreeNode) node).ne, x, y);
            if (leafOfNE != null) {
                return leafOfNE;
            }
            var leafOfSW = searchRecursively(((TreeNode) node).sw, x, y);
            if (leafOfSW != null) {
                return leafOfSW;
            }
            return searchRecursively(((TreeNode) node).se, x, y);
        } else {
            return null;
        }
    }

    private Leaf findLeafAt(int x, int y) {
        return searchRecursively(this.root, x, y);
    }

    public ArrayList<Point> findPath(Point startPoint, Point endPoint) {
        var startNode = findLeafAt(startPoint.x, startPoint.y);
        if (startNode == null || startNode.isObstacle) {
            return new ArrayList<Point>();
        }
        var endNode = findLeafAt(endPoint.x, endPoint.y);
        if (endNode == null || endNode.isObstacle) {
            return new ArrayList<Point>();
        }
        // 初始化开放集和关闭集,用于快速查找的节点映射。
        var openSet = new PriorityQueue<PathNode>();
        var closedSet = new HashSet<QuadTree>();
        var nodeMap = new HashMap<QuadTree, PathNode>();

        // 创建起点PathNode
        final var initialH = distance(startNode, endNode);
        final var initialNode = new PathNode(startNode, null, 0, initialH);

        openSet.add(initialNode);
        nodeMap.put(startNode, initialNode);

        // A*主循环
        boolean found = false;
        PathNode result = null;

        while (!openSet.isEmpty() && !found) {
            final var current = openSet.remove();
            if (current.treeNode == endNode) {
                found = true;
                result = current;
            }
            else
            {
                closedSet.add(current.treeNode);
                // 检查所有邻居。
                for (Leaf neighbor : current.treeNode.neighbors) {
                    if (closedSet.contains(neighbor)) {
                        continue;
                    }
                    final var tentativeG = current.g + distance(current.treeNode, neighbor);
                    var e = nodeMap.get(neighbor);
                    if (e == null || tentativeG < e.g) {
                        final var h = distance(neighbor, endNode);
                        var newNode = new PathNode(neighbor, current, tentativeG, h);
                        nodeMap.put(neighbor, newNode);
                        openSet.add(newNode);
                    }
                }
            }
        }

        if (!found) {
            return new ArrayList<>();
        }
        else
        {
            return reconstructPath(result);
        }
    }

    private ArrayList<Point> reconstructPath(PathNode result) {
        var l = new ArrayList<Point>();
        l.add(result.treeNode.center());
        var iter = result.parent;
        while(iter != null) {
            l.add(iter.treeNode.center());
            iter = iter.parent;
        }
        return l;
    }

}
